import cv2
import numpy as np
import os
import matplotlib.pyplot as plt

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False  


def process_hanzi_image(image_path):
    # 创建保存目录
    os.makedirs('chars', exist_ok=True)
    
    # 1. 读取原始图像（彩色图）
    img_color = cv2.imread(image_path)
    if img_color is None:
        print("无法加载图像，请检查路径")
        return
    # 转灰度图
    img_gray = cv2.cvtColor(img_color, cv2.COLOR_BGR2GRAY)

    # 2. 图像滤波（中值滤波，对应流程里的“图像滤波”）
    blurred = cv2.medianBlur(img_gray, 3)  # 3x3 中值滤波

    # 3. 全局阈值处理（二值化，对应“全局阈值处理<<二值化图”）
    _, binary = cv2.threshold(blurred, 127, 255, cv2.THRESH_BINARY_INV)

    # 4. 应用腐蚀操作（对应“应用腐蚀操作<<去除噪点”）
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    eroded = cv2.erode(binary, kernel, iterations=1)

    # 5. 应用膨胀操作（对应“应用膨胀操作<<突出图像特征”）
    dilated = cv2.dilate(eroded, kernel, iterations=1)

    # 6. 应用闭运算（对应“应用闭运算<<填充闭合区域”）
    closed = cv2.morphologyEx(dilated, cv2.MORPH_CLOSE, kernel, iterations=2)
    # 填充闭合区域（确保字符内部无空洞）
    filled = cv2.floodFill(closed, None, (0, 0), 255)[1]  # 从左上角开始填充

    # 7. Canny 边缘检测（对应“Canny边缘检测”）
    canny = cv2.Canny(filled, 50, 150)

    # 8. 查找轮廓（筛选汉字轮廓，关键步骤：找所有可能轮廓）
    contours, _ = cv2.findContours(filled, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    # 9. 筛选汉字轮廓（按面积、宽高比，适配字帖 100 个格子的特点）
    char_contours = []
    # 先统计所有轮廓的宽高，计算合理范围（适配统一格子）
    widths = []
    heights = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        widths.append(w)
        heights.append(h)
    # 用均值±标准差筛选（确保只保留汉字格子）
    if widths and heights:
        avg_w = np.mean(widths)
        std_w = np.std(widths)
        avg_h = np.mean(heights)
        std_h = np.std(heights)
        
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            area = cv2.contourArea(cnt)
            # 筛选条件：尺寸接近平均，面积合理（排除大轮廓如表格外框）
            if (avg_w - 2*std_w < w < avg_w + 2*std_w) and \
               (avg_h - 2*std_h < h < avg_h + 2*std_h) and \
               (area > avg_w*avg_h*0.5):  # 排除极小轮廓
                char_contours.append(cnt)

    # 10. 按坐标排序（确保从左到右、从上到下提取，对应字帖顺序）
    char_contours.sort(key=lambda cnt: (cv2.boundingRect(cnt)[1], cv2.boundingRect(cnt)[0]))

    # 11. 提取并保存单个汉字（精准切割，保证 100 个）
    valid_count = 0
    for i, cnt in enumerate(char_contours):
        x, y, w, h = cv2.boundingRect(cnt)
        # 切割汉字区域（无额外 padding，严格切割）
        char_img = img_gray[y:y+h, x:x+w]
        # 二值化确保黑字白底
        _, char_binary = cv2.threshold(char_img, 127, 255, cv2.THRESH_BINARY_INV)
        # 保存（命名格式和目标一致）
        cv2.imwrite(f'chars/char_{i:03d}.png', char_binary)
        valid_count += 1
        # 提前终止（如果已知共 100 个，避免多余轮廓干扰）
        if valid_count == 100:
            break

    # 可视化处理流程（严格对应你提供的步骤）
    plt.figure(figsize=(12, 8))
    # 1. 原图
    plt.subplot(2, 4, 1)
    plt.imshow(cv2.cvtColor(img_color, cv2.COLOR_BGR2RGB))
    plt.title('1. 原图')
    plt.axis('off')
    # 2. 灰度图
    plt.subplot(2, 4, 2)
    plt.imshow(img_gray, cmap='gray')
    plt.title('2. 灰度图')
    plt.axis('off')
    # 3. 滤波后
    plt.subplot(2, 4, 3)
    plt.imshow(blurred, cmap='gray')
    plt.title('3. 图像滤波（中值）')
    plt.axis('off')
    # 4. 二值化
    plt.subplot(2, 4, 4)
    plt.imshow(binary, cmap='gray')
    plt.title('4. 全局阈值二值化')
    plt.axis('off')
    # 5. 腐蚀后
    plt.subplot(2, 4, 5)
    plt.imshow(eroded, cmap='gray')
    plt.title('5. 应用腐蚀操作')
    plt.axis('off')
    # 6. 膨胀后
    plt.subplot(2, 4, 6)
    plt.imshow(dilated, cmap='gray')
    plt.title('6. 应用膨胀操作')
    plt.axis('off')
    # 7. 闭运算+填充后
    plt.subplot(2, 4, 7)
    plt.imshow(filled, cmap='gray')
    plt.title('7. 应用闭运算+填充')
    plt.axis('off')
    # 8. Canny边缘
    plt.subplot(2, 4, 8)
    plt.imshow(canny, cmap='gray')
    plt.title('8. Canny边缘检测')
    plt.axis('off')
    plt.tight_layout()
    plt.show()

    print(f"成功提取 {valid_count} 个汉字到 chars 目录")


# 调用示例（替换为你的字帖图片路径）
process_hanzi_image('hanzi1.jpg')